package com.academic.web.controller.scala;

import com.academic.system.domain.Student;
import lombok.extern.slf4j.Slf4j;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.sql.*;
import org.apache.spark.ml.feature.VectorAssembler;
import org.apache.spark.ml.regression.LinearRegression;
import org.apache.spark.ml.regression.LinearRegressionModel;
import org.apache.spark.ml.regression.LinearRegressionTrainingSummary;
import org.apache.spark.ml.evaluation.RegressionEvaluator;
import org.apache.spark.sql.types.DataTypes;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.List;

import static org.apache.spark.sql.functions.*;

@Slf4j
@Component
public class AcademicProcess {
    @Value("${hdfs.host}")
    private String host;

    @Value("${hdfs.uploadPath}")
    private String uploadPath;

    public  List<Student> getWarn(String hdfspath) {
        // 调整Spark配置
        SparkConf conf = new SparkConf()
                .setAppName("AcademicProcess")
                .setMaster("local[*]")
                .set("spark.executor.memory", "2g")
                .set("spark.driver.memory", "2g")
                .set("spark.sql.shuffle.partitions", "5");

        JavaSparkContext sc = new JavaSparkContext(conf);
        SparkSession spark = SparkSession.builder().config(conf).getOrCreate();

         hdfspath = host+hdfspath;

        // 读取数据
        Dataset<Row> data = spark.read()
                .format("com.crealytics.spark.excel")
                .option("header", "true")
                .option("inferSchema", "true")
                .option("dataAddress", "A3")
                .load(hdfspath);

        log.info("原始数据:");
        data.show();

        data = data.withColumn("总成绩", data.col("总成绩").cast(DataTypes.DoubleType))
                .withColumn("学时", data.col("学时").cast(DataTypes.DoubleType))
                .withColumn("学分", data.col("学分").cast(DataTypes.DoubleType))
                .na().fill(0.0, new String[]{"总成绩", "学时", "学分"});

        log.info("处理后的数据:");
        data.show();
        VectorAssembler assembler = new VectorAssembler()
                .setInputCols(new String[]{"总成绩", "学时", "学分"})
                .setOutputCol("features");

        Dataset<Row> assembledData = assembler.transform(data);

        Dataset<Row>[] splits = assembledData.randomSplit(new double[]{0.7, 0.3});
        Dataset<Row> trainingData = splits[0];
        Dataset<Row> testData = splits[1];

        LinearRegression lr = new LinearRegression()
                .setLabelCol("总成绩")
                .setFeaturesCol("features");

        LinearRegressionModel model = lr.fit(trainingData);

        LinearRegressionTrainingSummary summary = model.summary();
        log.info("RMSE: " + summary.rootMeanSquaredError());
        log.info("r2: " + summary.r2());

        Dataset<Row> predictions = model.transform(testData);
        log.info("预测结果:");
        predictions.show();

        Dataset<Row> lowScoreStudents = predictions.filter(predictions.col("prediction").lt(60));

        Dataset<Row> alertCounts = lowScoreStudents.groupBy("学号", "姓名", "上课院系")
                .agg(count("学号").alias("预警次数"));

        Dataset<Row> alertStudents = alertCounts.withColumn("预警等级",
                when(col("预警次数").equalTo(1), "一级预警")
                        .when(col("预警次数").equalTo(2), "二级预警")
                        .otherwise("三级预警")
        );

        log.info("预警学生的学号、姓名、学院和预警等级：");
        alertStudents.select("学号", "姓名", "上课院系", "预警等级").show();

        RegressionEvaluator evaluator = new RegressionEvaluator()
                .setLabelCol("总成绩")
                .setPredictionCol("prediction")
                .setMetricName("rmse");

        double rmse = evaluator.evaluate(predictions);
        log.info("Root Mean Squared Error (RMSE) on test data = " + rmse);

        // 将列名重命名为与 Student 类属性名一致
        Dataset<Row> renamedAlertStudents = alertStudents
                .withColumnRenamed("学号", "id")
                .withColumnRenamed("姓名", "name")
                .withColumnRenamed("上课院系", "department")
                .withColumnRenamed("预警等级", "alertLevel");

        // 将数据集转换为 List<Student>
        List<Student> studentsList = renamedAlertStudents.select("id", "name", "department", "alertLevel")
                .as(Encoders.bean(Student.class))
                .collectAsList();

        sc.stop();

        return studentsList;
    }

    public static void main(String[] args) {
        AcademicProcess academicProcess = new AcademicProcess();
        List<Student> warn = academicProcess.getWarn("/upload/xs.xls");
        System.out.println(warn);
//        List<Student> warn = getWarn();
//        for (Student student : warn) {
//            System.out.println(student);
//        }
    }
}
